
import visualize
from torchvision.ops import boxes
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import subprocess
import os
import urllib.parse
import psutil
from numpy import asarray
from shapely.geometry import Polygon


def find_points(pts):

    rec = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rec[0] = pts[np.argmin(s)]
    rec[2] = pts[np.argmax(s)]

    diff = np.diff(pts, axis=1)
    rec[1] = pts[np.argmin(diff)]
    rec[3] = pts[np.argmax(diff)]

    return rec

# calculates iou between two polygons

def IOU(box_1, box_2):
    poly_1 = Polygon(box_1)
    poly_2 = Polygon(box_2)
    iou = poly_1.intersection(poly_2).area / poly_1.union(poly_2).area
    return iou


def detect_corners(image):
    # YOLO model trained to detect corners on a chessboard
    model_trained = YOLO("corners_recognise.pt")
    results = model_trained.predict(source=image, line_width=1, conf=0.25, save_txt=True, save=True)

    # get the corners coordinates from the model
    boxes = results[0].boxes
    arr = boxes.xywh.cpu().numpy()
    points = arr[:, 0:2]

    corners = find_points(points)

    return corners


def four_point(image, pts, top_offset=50):
  
    img = Image.open(image)
    image = asarray(img)
    # 复制点坐标，避免修改原图
    adjusted_pts = pts.copy()
    # 调整左上和右上点的Y坐标，使其向上移动
    if top_offset != 0:
        adjusted_pts[0][1] -= top_offset  # 左上点Y坐标减小
        adjusted_pts[1][1] -= top_offset  # 右上点Y坐标减小
    rec = find_points(adjusted_pts)
    (tl, tr, br, bl) = rec
    # 计算新图像的宽度
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))
    # 计算新图像的高度
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))
    # 构造目标点以获得鸟瞰图
    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype="float32")
    # 计算透视变换矩阵并应用变换
    M = cv2.getPerspectiveTransform(rec, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
    img = Image.fromarray(warped, "RGB")
    return img




def find_points(pts):
    """对输入的四个点进行排序，顺序为：左上、右上、右下、左下"""
    rec = np.zeros((4, 2), dtype="float32")
    # 计算点的x和y坐标之和
    s = pts.sum(axis=1)
    # 左上点的和最小，右下点的和最大
    rec[0] = pts[np.argmin(s)]
    rec[2] = pts[np.argmax(s)]

    # 计算点的x和y坐标之差
    diff = np.diff(pts, axis=1)
    # 右上点的差最小，左下点的差最大
    rec[1] = pts[np.argmin(diff)]
    rec[3] = pts[np.argmax(diff)]

    return rec


def grid_image(image):
    # 确保输入是OpenCV格式（BGR）
    if isinstance(image, Image.Image):
        img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    else:
        img = image.copy()

    # 获取图像尺寸
    h, w = img.shape[:2]

    # 定义四个角点（顺序：左上、右上、右下、左下）
    corners = np.array([[0, 0], [w, 0], [w, h], [0, h]], dtype=np.float32)
    corners = find_points(corners)

    # 提取角点坐标
    TL = corners[0]  # 左上
    TR = corners[1]  # 右上
    BR = corners[2]  # 右下
    BL = corners[3]  # 左下

    # 定义插值函数（生成网格点）
    def interpolate(xy0, xy1, n=8):
        """在两点之间插值生成n+1个点"""
        x0, y0 = xy0
        x1, y1 = xy1
        return [(x0 + i * (x1 - x0) / n, y0 + i * (y1 - y0) / n) for i in range(n + 1)]

    # 计算网格线的坐标点
    ptsT = interpolate(TL, TR)  # 顶部边缘的点
    ptsL = interpolate(TL, BL)  # 左侧边缘的点
    ptsR = interpolate(TR, BR)  # 右侧边缘的点
    ptsB = interpolate(BL, BR)  # 底部边缘的点

    # 绘制网格线
    color = (0, 0, 255)  # 红色 (BGR格式)
    thickness = 1

    # 绘制垂直线
    for i in range(9):
        pt1 = (int(ptsL[i][0]), int(ptsL[i][1]))
        pt2 = (int(ptsR[i][0]), int(ptsR[i][1]))
        cv2.line(img, pt1, pt2, color, thickness, cv2.LINE_AA)

    # 绘制水平线
    for i in range(9):
        pt1 = (int(ptsT[i][0]), int(ptsT[i][1]))
        pt2 = (int(ptsB[i][0]), int(ptsB[i][1]))
        cv2.line(img, pt1, pt2, color, thickness, cv2.LINE_AA)

    # 保存图像
    cv2.imwrite('chessboard_transformed_with_grid.jpg', img)

    # 返回网格点坐标（与原代码保持一致）
    return ptsT, ptsL


# detects chess pieces

def chess_detects(image):
    model_trained = YOLO("chess_recognise.pt")
    results = model_trained.predict(source=image, line_width=1, conf=0.5, augment=False, save_txt=True, save=True)

    boxes = results[0].boxes
    detections = boxes.xyxy.cpu().numpy()

    return detections, boxes


# connects detected piece to the right square

def connect_to_detection(detections, square):
    di = {0: 'b', 1: 'k', 2: 'n',
          3: 'p', 4: 'q', 5: 'r',
          6: 'B', 7: 'K', 8: 'N',
          9: 'P', 10: 'Q', 11: 'R'}

    list_of_iou = []

    for i in detections:

        box_x1 = i[0]
        box_y1 = i[1]

        box_x2 = i[2]
        box_y2 = i[1]

        box_x3 = i[2]
        box_y3 = i[3]

        box_x4 = i[0]
        box_y4 = i[3]

        # cut high pieces
        if box_y4 - box_y1 > 60:
            box_complete = np.array([[box_x1, box_y1 + 40], [box_x2, box_y2 + 40], [box_x3, box_y3], [box_x4, box_y4]])
        else:
            box_complete = np.array([[box_x1, box_y1], [box_x2, box_y2], [box_x3, box_y3], [box_x4, box_y4]])

        # until here

        list_of_iou.append(IOU(box_complete, square))

    num = list_of_iou.index(max(list_of_iou))

    piece = boxes.cls[num].tolist()

    if max(list_of_iou) > 0.15:
        piece = boxes.cls[num].tolist()
        return di[piece]

    else:
        piece = "empty"
        return piece


def is_arena_running():
    """检查Arena进程是否正在运行"""
    for proc in psutil.process_iter(['name']):
        if 'Arena.exe' in proc.info['name']:
            return True
    return False


def launch_arena_with_fen(arena_path, fen):
    """
    使用指定的FEN码启动或控制Arena国际象棋界面

    参数:
    arena_path (str): Arena可执行文件的路径
    fen (str): 有效的FEN棋盘表示法
    """
    # 检查Arena可执行文件是否存在
    if not os.path.exists(arena_path):
        print(f"错误: Arena可执行文件不存在于路径 '{arena_path}'")
        return False

    # 对FEN进行URL编码，因为它将作为命令行参数传递
    encoded_fen = urllib.parse.quote(fen)

    try:
        if not is_arena_running():
            # 如果Arena未运行，则启动它
            command = [arena_path, "/FEN", encoded_fen]
            subprocess.Popen(command)
            print(f"Arena已启动，初始棋盘布局为: {fen}")
        else:
            # 如果Arena已运行，尝试通过发送FEN码来改变棋盘
            # 注意：这种方法可能需要根据Arena的具体实现进行调整
            # 这里使用Windows的WM_COPYDATA消息来发送FEN码
            import win32gui
            import win32con
            import struct

            # 查找Arena窗口
            arena_window = win32gui.FindWindow(None, "Arena Chess")
            if arena_window == 0:
                # 尝试使用可能的其他窗口标题
                arena_window = win32gui.FindWindow(None, "Arena")

            if arena_window != 0:
                # 准备要发送的数据
                # 格式: [消息类型(1=设置FEN), FEN字符串]
                message_type = 1
                data = f"{message_type}|{fen}".encode('utf-8')

                # 发送WM_COPYDATA消息
                cds = struct.pack("LLP", win32con.WM_COPYDATA, len(data), data)
                win32gui.SendMessage(arena_window, win32con.WM_COPYDATA, 0, cds)

                print(f"Arena已在运行，棋盘布局已更新为: {fen}")
            else:
                print("错误: 找不到正在运行的Arena窗口")
                return False

        return True
    except Exception as e:
        print(f"操作Arena时出错: {str(e)}")
        return False


def validate_fen(fen):
    """验证FEN格式是否正确"""
    # 这里是简化的验证，实际应用中可能需要更严格的检查
    parts = fen.strip().split()
    if len(parts) != 6:
        return False

    board, turn, castling, en_passant, halfmove, fullmove = parts

    # 简单检查棋盘部分
    rows = board.split('/')
    if len(rows) != 8:
        return False

    # 检查回合部分
    if turn not in ('w', 'b'):
        return False

    return True


def fix_fen(fen: str) -> str:
    """
    修复FEN码中的数字格式问题并补全缺失部分

    Args:
        fen: 待修复的FEN码字符串

    Returns:
        修复后的完整FEN码
    """
    try:
        # 分割棋盘布局与其他部分
        parts = fen.strip().split()

        if not parts:
            raise ValueError("空的FEN码")

        # 处理棋盘布局部分
        board_part = parts[0]
        ranks = board_part.split('/')

        if len(ranks) != 8:
            raise ValueError(f"棋盘必须有8行，但得到了{len(ranks)}行")

        fixed_ranks = []
        for rank in ranks:
            # 合并连续数字并确保每行总长度为8
            new_rank = ""
            current_num = 0

            for c in rank:
                if c.isdigit():
                    current_num += int(c)
                else:
                    if current_num > 0:
                        new_rank += str(current_num)
                        current_num = 0
                    new_rank += c

            # 添加剩余数字
            if current_num > 0:
                new_rank += str(current_num)

            # 修正超过8的数字
            final_rank = ""
            remaining = 8
            for c in new_rank:
                if c.isdigit():
                    num = int(c)
                    if num > remaining:
                        final_rank += str(remaining)
                        remaining = 0
                    else:
                        final_rank += c
                        remaining -= num
                else:
                    final_rank += c
                    remaining -= 1

            # 检查修正后的行长度
            if remaining > 0:
                final_rank += str(remaining)

            fixed_ranks.append(final_rank)

        # 构建修正后的棋盘布局
        fixed_board = '/'.join(fixed_ranks)

        # 补全缺失的FEN部分
        if len(parts) == 1:
            # 添加默认的其他部分: 白方回合, 双方有易位权, 无吃过路兵, 半回合0, 总回合1
            return f"{fixed_board} w KQkq - 0 1"
        else:
            return f"{fixed_board} {' '.join(parts[1:])}"

    except Exception as e:
        print(f"FEN码修复失败: {e}")
        return None

if __name__ == '__main__':
    image = "pic3.jpg"

    corners = detect_corners(image)

    transformed_image = four_point(image, corners)

    ptsT, ptsL = grid_image(transformed_image)

    detections, boxes = chess_detects(transformed_image)

    # calculate the grid

    x_coords = [ptsT[i][0] for i in range(9)]  # 包含x0到x8（对应原xA到xI）
    y_coords = [ptsL[i][1] for i in range(9)]  # 包含y0到y8（对应原y9到y1，索引反转可对应棋盘行）

    # 2. 动态生成所有棋盘格子（8x8）
    # 注意：国际象棋棋盘行号从8到1，对应代码中row=0到7（a8对应row=0, col=0）
    chessboard_squares = []
    for row in range(8):
        row_squares = []
        for col in range(8):
            # 计算当前格子的四个角点坐标（顺序：左上、右上、右下、左下）
            x1, x2 = x_coords[col], x_coords[col + 1]
            y1, y2 = y_coords[row], y_coords[row + 1]
            square = np.array([
                [x1, y1],  # 左上
                [x2, y1],  # 右上
                [x2, y2],  # 右下
                [x1, y2]  # 左下
            ])
            row_squares.append(square)
        chessboard_squares.append(row_squares)
    # transforms the squares to write FEN

    FEN_annotation = chessboard_squares

    board_FEN = []
    correced_FEN = []
    complete_board_FEN = []
    for line in FEN_annotation:
        line_to_FEN = []
        for square in line:
            piece_on_square = connect_to_detection(detections, square)
            line_to_FEN.append(piece_on_square)
        correced_FEN = [i.replace('empty', '1') for i in line_to_FEN]
        print(correced_FEN)
        board_FEN.append(correced_FEN)

    complete_board_FEN = [''.join(line) for line in board_FEN]

    to_FEN = '/'.join(complete_board_FEN)
    FEN=fix_fen(to_FEN)
    print(FEN)

    root = visualize.tk.Tk()
    app = visualize.ChessGUI(root)
    app.load_fen(FEN)
    # 传入默认的fen码
    root.mainloop()